{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "14d8e34a-cb69-4459-a848-4002ae004365",
   "metadata": {},
   "source": [
    "### MNIST数据集\n",
    "- 60000张训练图片\n",
    "- 10000张测试图片\n",
    "- 数据集为灰度图\n",
    "- 图像大小为28 * 28\n",
    "\n",
    "### pytorch搭建神经网络的一般步骤：\n",
    "- 加载数据\n",
    "- 定义网络模型\n",
    "- 损失函数\n",
    "- 优化器\n",
    "- 训练\n",
    "- 测试\n",
    "- 保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6e3df5b0-b068-4d5b-ba16-7f4eed4517f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导包\n",
    "import torch\n",
    "from torchvision import datasets\n",
    "from torchvision import transforms\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c42e10b-1e43-4524-a914-4e9647ceed1b",
   "metadata": {},
   "source": [
    "#### 加载数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92ec827b-90e8-41af-a72c-05412bc8729b",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = datasets.MNIST(root='data/mnist', train=True, transform=transforms.ToTensor(), download=True)\n",
    "test_data = datasets.MNIST(root='data/mnist', train=False, transform=transforms.ToTensor(), download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ac73785e-23df-46b9-97fa-8e9b2eea89ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 100 # 每次选取一小部分数据进行训练\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89225480-71cd-4bf7-9eeb-f9a13ebdbcd9",
   "metadata": {},
   "source": [
    "#### 网络构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "203894a1-3e40-4d34-a5dc-9a0b6d8bf977",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义MLP网络\n",
    "class MLP(nn.Module):\n",
    "\n",
    "    # 初始化方法\n",
    "    # input_size：输入数据的维度\n",
    "    # hidden_size: 隐藏层的大小\n",
    "    # num_classes: 输出分类的数量\n",
    "    def __init__(self, input_size, hidden_size, num_classes):\n",
    "        super(MLP, self).__init__()\n",
    "        # 定义第一个全连接层\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        # 定义激活函数\n",
    "        self.relu = nn.ReLU()\n",
    "        # 定义第二个全连接层\n",
    "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
    "        # 定义第三个全连接层\n",
    "        self.fc3 = nn.Linear(hidden_size, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.fc1(x)\n",
    "        out = self.relu(out)\n",
    "        out = self.fc2(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.fc3(out)\n",
    "        return out\n",
    "\n",
    "input_size = 28 * 28 # 输入大小\n",
    "hidden_size = 512 # 隐藏层大小\n",
    "num_classes = 10 # 输出类别数\n",
    "\n",
    "# 初始化MLP\n",
    "model = MLP(input_size, hidden_size, num_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f06693b-24c6-4578-8f41-265aa2ae59c4",
   "metadata": {},
   "source": [
    "#### Loss\n",
    "因为是分类问题，所以使用交叉熵损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "21982269-b4f5-443d-8383-7be05e8143ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5cf1084-0234-44a5-80e5-9cc7806fec10",
   "metadata": {},
   "source": [
    "#### 优化器optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "024ba02a-22f4-4f2c-9708-674c5852ca7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.001 # 学习率\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e807dc8b-cd1a-4929-a4e1-143b3ea56401",
   "metadata": {},
   "source": [
    "#### 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ec9a3386-847c-4c29-9766-9ec0341c1d45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1 / 10], Step [100 / 600], Loss:  0.3921\n",
      "Epoch [1 / 10], Step [200 / 600], Loss:  0.2537\n",
      "Epoch [1 / 10], Step [300 / 600], Loss:  0.1324\n",
      "Epoch [1 / 10], Step [400 / 600], Loss:  0.0847\n",
      "Epoch [1 / 10], Step [500 / 600], Loss:  0.1240\n",
      "Epoch [1 / 10], Step [600 / 600], Loss:  0.1721\n",
      "Epoch [2 / 10], Step [100 / 600], Loss:  0.0946\n",
      "Epoch [2 / 10], Step [200 / 600], Loss:  0.0753\n",
      "Epoch [2 / 10], Step [300 / 600], Loss:  0.1020\n",
      "Epoch [2 / 10], Step [400 / 600], Loss:  0.0737\n",
      "Epoch [2 / 10], Step [500 / 600], Loss:  0.1320\n",
      "Epoch [2 / 10], Step [600 / 600], Loss:  0.0956\n",
      "Epoch [3 / 10], Step [100 / 600], Loss:  0.0219\n",
      "Epoch [3 / 10], Step [200 / 600], Loss:  0.0785\n",
      "Epoch [3 / 10], Step [300 / 600], Loss:  0.0193\n",
      "Epoch [3 / 10], Step [400 / 600], Loss:  0.0106\n",
      "Epoch [3 / 10], Step [500 / 600], Loss:  0.0211\n",
      "Epoch [3 / 10], Step [600 / 600], Loss:  0.0655\n",
      "Epoch [4 / 10], Step [100 / 600], Loss:  0.0454\n",
      "Epoch [4 / 10], Step [200 / 600], Loss:  0.0524\n",
      "Epoch [4 / 10], Step [300 / 600], Loss:  0.0447\n",
      "Epoch [4 / 10], Step [400 / 600], Loss:  0.0661\n",
      "Epoch [4 / 10], Step [500 / 600], Loss:  0.0416\n",
      "Epoch [4 / 10], Step [600 / 600], Loss:  0.0081\n",
      "Epoch [5 / 10], Step [100 / 600], Loss:  0.0745\n",
      "Epoch [5 / 10], Step [200 / 600], Loss:  0.0708\n",
      "Epoch [5 / 10], Step [300 / 600], Loss:  0.0660\n",
      "Epoch [5 / 10], Step [400 / 600], Loss:  0.0728\n",
      "Epoch [5 / 10], Step [500 / 600], Loss:  0.0091\n",
      "Epoch [5 / 10], Step [600 / 600], Loss:  0.0073\n",
      "Epoch [6 / 10], Step [100 / 600], Loss:  0.0018\n",
      "Epoch [6 / 10], Step [200 / 600], Loss:  0.0367\n",
      "Epoch [6 / 10], Step [300 / 600], Loss:  0.0281\n",
      "Epoch [6 / 10], Step [400 / 600], Loss:  0.0332\n",
      "Epoch [6 / 10], Step [500 / 600], Loss:  0.0021\n",
      "Epoch [6 / 10], Step [600 / 600], Loss:  0.0430\n",
      "Epoch [7 / 10], Step [100 / 600], Loss:  0.0026\n",
      "Epoch [7 / 10], Step [200 / 600], Loss:  0.0167\n",
      "Epoch [7 / 10], Step [300 / 600], Loss:  0.0212\n",
      "Epoch [7 / 10], Step [400 / 600], Loss:  0.0250\n",
      "Epoch [7 / 10], Step [500 / 600], Loss:  0.0049\n",
      "Epoch [7 / 10], Step [600 / 600], Loss:  0.0110\n",
      "Epoch [8 / 10], Step [100 / 600], Loss:  0.0019\n",
      "Epoch [8 / 10], Step [200 / 600], Loss:  0.0232\n",
      "Epoch [8 / 10], Step [300 / 600], Loss:  0.0010\n",
      "Epoch [8 / 10], Step [400 / 600], Loss:  0.0300\n",
      "Epoch [8 / 10], Step [500 / 600], Loss:  0.0158\n",
      "Epoch [8 / 10], Step [600 / 600], Loss:  0.0047\n",
      "Epoch [9 / 10], Step [100 / 600], Loss:  0.0013\n",
      "Epoch [9 / 10], Step [200 / 600], Loss:  0.0011\n",
      "Epoch [9 / 10], Step [300 / 600], Loss:  0.0076\n",
      "Epoch [9 / 10], Step [400 / 600], Loss:  0.0005\n",
      "Epoch [9 / 10], Step [500 / 600], Loss:  0.0120\n",
      "Epoch [9 / 10], Step [600 / 600], Loss:  0.0024\n",
      "Epoch [10 / 10], Step [100 / 600], Loss:  0.0010\n",
      "Epoch [10 / 10], Step [200 / 600], Loss:  0.0023\n",
      "Epoch [10 / 10], Step [300 / 600], Loss:  0.0012\n",
      "Epoch [10 / 10], Step [400 / 600], Loss:  0.0077\n",
      "Epoch [10 / 10], Step [500 / 600], Loss:  0.0151\n",
      "Epoch [10 / 10], Step [600 / 600], Loss:  0.0355\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 10 # 训练轮数,训练10次\n",
    "for epoch in range(num_epochs):\n",
    "    for i, (images, labels) in enumerate(train_loader): # 每次读取batch_size大小的数据\n",
    "        # 将image转换为向量\n",
    "        images = images.reshape(-1, 28 * 28)\n",
    "        # 将数据输入到网络中\n",
    "        outputs = model(images)\n",
    "        # 计算损失\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        # 首先将梯度清零\n",
    "        optimizer.zero_grad()\n",
    "        # 反向传播\n",
    "        loss.backward()\n",
    "        # 更新参数\n",
    "        optimizer.step()\n",
    "\n",
    "        if (i + 1) % 100 == 0:\n",
    "            print(f'Epoch [{epoch + 1} / {num_epochs}], Step [{i + 1} / {len(train_loader)}], Loss: {loss.item(): .4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01437c01-a653-4a4b-a527-f833be56ecda",
   "metadata": {},
   "source": [
    "#### test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "871e83eb-7a08-49d8-a92e-7057501dfb48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy of the network on the 10000 test images: 97.86 %\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    # 从test_loader循环读取测试数据\n",
    "    for images, labels in test_loader:\n",
    "        # 将image转换为向量\n",
    "        images = images.reshape(-1, 28 * 28)\n",
    "        # 将数据输入网络\n",
    "        outputs = model(images)\n",
    "        # 取出最大值对应的索引，也就是预测值\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        # 累加label数\n",
    "        total += labels.size(0)\n",
    "        # 预测值与labels值比对 获取正确的数量\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "    # 打印最终的准确率\n",
    "    print(f'Accuracy of the network on the 10000 test images: {100 * correct / total} %')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b656c9c7-398e-4d8d-9525-f1c34647df64",
   "metadata": {},
   "source": [
    "#### save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cde51096-3073-4ed6-8a79-845d200854a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model, \"mnist_mlp_model.pkl\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
